#include <Eigen/Dense>
#include <fstream>
#include <iostream>
#include <random>

// 定义状态向量和测量向量
using State = Eigen::VectorXd;
using Measurement = Eigen::VectorXd;

// 定义噪声参数
const double Q_pos = 1.0;           // 位置噪声方差
const double Q_att = 0.002;         // 姿态噪声方差
const double Q_meas_alpha = 0.002;  // 俯仰角测量噪声方差
const double Q_meas_beta = 0.002;   // 方位角测量噪声方差
const double Q_meas_L = 2.0;        // 测距噪声方差

// 定义仿真参数
const double dt = 0.005;    // 仿真步长
const int step_num = 1000;  // 仿真步数

// 状态转移函数
State state_transition(const State& state) {
  State next_state = state;
  // 假设目标做匀速运动
  next_state(0) += state(3) * dt;
  next_state(1) += state(4) * dt;
  next_state(2) += state(5) * dt;
  return next_state;
}

// 观测模型
Measurement measurement_model(const State& state) {
  Measurement measurement(3);
  // 假设无人机静止，目标做匀速运动
  double range =
      sqrt(state(0) * state(0) + state(1) * state(1) + state(2) * state(2));
  double bearing_elevation = asin(state(1) / range);
  double bearing_azimuth = -atan2(state(2), state(0));

  // 将测量值转换为俯仰角和方位角
  measurement(0) = bearing_elevation;
  measurement(1) = bearing_azimuth;
  measurement(2) = range;

  return measurement;
}

// 添加测量噪声
void add_measurement_noise(Measurement& measurement) {
  std::default_random_engine generator;
  std::normal_distribution<double> noise_alpha(0, sqrt(Q_meas_alpha));
  std::normal_distribution<double> noise_beta(0, sqrt(Q_meas_beta));
  std::normal_distribution<double> noise_L(0, sqrt(Q_meas_L));

  measurement(0) += noise_alpha(generator);
  measurement(1) += noise_beta(generator);
  measurement(2) += noise_L(generator);
}

// 雅可比矩阵
Eigen::MatrixXd Jacobian(const State& state) {
  Eigen::MatrixXd H(3, 6);
  double x = state(0);
  double y = state(1);
  double z = state(2);
  double range = sqrt(x * x + y * y + z * z);

  // 俯仰角的雅可比矩阵
  H(0, 0) = (y * x) / (x * x + y * y + z * z) * (1 / (range * range));
  H(0, 1) =
      (y * y) / (x * x + y * y + z * z) * (1 / (range * range)) + (1 / range);
  H(0, 2) = (y * z) / (x * x + y * y + z * z) * (1 / (range * range));
  H(0, 3) = 0;
  H(0, 4) = 0;
  H(0, 5) = 0;

  // 方位角的雅可比矩阵
  H(1, 0) = -z / (x * sqrt(x * x + z * z));  // 相对于x的偏导数
  H(1, 1) = 0;                               // 相对于y的偏导数
  H(1, 2) = x / (x * sqrt(x * x + z * z));   // 相对于z的偏导数
  H(1, 3) = 0;                               // 相对于dxdot的偏导数
  H(1, 4) = 0;                               // 相对于dydot的偏导数
  H(1, 5) = 0;                               // 相对于dzdot的偏导数

  // 距离的雅可比矩阵
  H(2, 0) = x / range;  // 相对于x的偏导数
  H(2, 1) = y / range;  // 相对于y的偏导数
  H(2, 2) = z / range;  // 相对于z的偏导数
  H(2, 3) = dt;         // 相对于dxdot的偏导数
  H(2, 4) = dt;         // 相对于dydot的偏导数
  H(2, 5) = dt;         // 相对于dzdot的偏导数

  return H;
}

// 预测步骤
State predict(const State& state, const Eigen::MatrixXd& P,
              const Eigen::MatrixXd& Q, const Eigen::MatrixXd& F) {
  State predicted_state = state_transition(state);
  Eigen::MatrixXd predicted_P = F * P * F.transpose() + Q;
  return predicted_state;
}

// 更新步骤
State update(const State& state, const Eigen::MatrixXd& P,
             const Eigen::MatrixXd& H, const Measurement& measurement,
             const Eigen::MatrixXd& R) {
  Eigen::MatrixXd K = P * H.transpose() * (H * P * H.transpose() + R).inverse();
  State updated_state = state + K * (measurement - measurement_model(state));
  Eigen::MatrixXd updated_P = (Eigen::MatrixXd::Identity(6, 6) - K * H) * P;
  return updated_state;
}

int main() {
  // 初始化状态和协方差矩阵
  State state(6);
  state << 100, -100, 0, 2, 0, 0;  // 目标初始位置相对于无人机的位置
  Eigen::MatrixXd P = Eigen::MatrixXd::Identity(6, 6) * 0.5;  //协方差矩阵
  Eigen::MatrixXd Q(6, 6);
  Q << Q_pos, 0, 0, 0, 0, 0, 0, Q_pos, 0, 0, 0, 0, 0, 0, Q_pos, 0, 0, 0, 0, 0,
      0, Q_pos, 0, 0, 0, 0, 0, 0, Q_pos, 0, 0, 0, 0, 0, 0, Q_pos;
  Eigen::MatrixXd R(3, 3);
  R << Q_meas_alpha, 0, 0, 0, Q_meas_beta, 0, 0, 0, Q_meas_L;
  Eigen::MatrixXd F(6, 6);
  F << 1, 0, 0, dt, 0, 0, 0, 1, 0, 0, dt, 0, 0, 0, 1, 0, 0, dt, 0, 0, 0, 1, 0,
      0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1;

  // 打开文件用于写入
  std::ofstream outFile("simulation_data.txt");

  // 目标初始位置和速度
  State true_state(6);
  true_state << 100, -100, 0, 2, 0, 0;

  for (int step = 0; step < step_num; ++step) {
    // 预测步骤
    State predicted_state = predict(state, P, Q, F);

    // 计算真实测量量
    Measurement true_measurement = measurement_model(true_state);
    add_measurement_noise(true_measurement);

    // 更新步骤
    Eigen::MatrixXd H = Jacobian(state);
    state = update(state, P, H, true_measurement, R);

    // 输出状态和真值到文件
    outFile << step << " " << state(0) << " " << state(1) << " " << state(2)
            << " " << state(3) << " " << state(4) << " " << state(5) << " "
            << true_state(0) << " " << true_state(1) << " " << true_state(2)
            << " " << true_state(3) << " " << true_state(4) << " "
            << true_state(5) << std::endl;
    // 为下一次迭代更新真实状态
    true_state(0) += 2 * dt;  // 目标以2m/s的速度在x轴方向移动
  }

  // 关闭文件
  outFile.close();

  return 0;
}